"""
三维电磁场预测模型 - 基于Fourier神经算子(FNO)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class SpectralConv3d(nn.Module):
    """3D傅里叶卷积层，用于处理频域信息"""
    
    def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
        super(SpectralConv3d, self).__init__()
        """初始化3D傅里叶卷积层
        
        参数:
            in_channels: 输入通道数
            out_channels: 输出通道数
            modes1: x方向傅里叶模式数
            modes2: y方向傅里叶模式数
            modes3: z方向傅里叶模式数
        """
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1
        self.modes2 = modes2
        self.modes3 = modes3
        
        # 可学习的傅里叶系数
        scale = 1 / (in_channels * out_channels)
        self.weights1 = nn.Parameter(
            scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)
        )
        self.weights2 = nn.Parameter(
            scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)
        )
    
    def forward(self, x):
        """前向传播
        
        参数:
            x: 输入张量，形状为 [batch_size, in_channels, x_dim, y_dim, z_dim]
            
        返回:
            输出张量，形状为 [batch_size, out_channels, x_dim, y_dim, z_dim]
        """
        batchsize = x.shape[0]
        
        # 计算傅里叶变换
        x_ft = torch.fft.rfftn(x, dim=[-3, -2, -1])
        
        # 处理低频部分
        out_ft = torch.zeros(
            batchsize, self.out_channels, x.size(-3), x.size(-2), x.size(-1)//2 + 1,
            dtype=torch.cfloat, device=x.device
        )
        
        # 与可学习权重相乘
        out_ft[:, :, :self.modes1, :self.modes2, :self.modes3] = \
            self.complex_mul3d(x_ft[:, :, :self.modes1, :self.modes2, :self.modes3], self.weights1)
        
        # 处理高频部分
        out_ft[:, :, -self.modes1:, :self.modes2, :self.modes3] = \
            self.complex_mul3d(x_ft[:, :, -self.modes1:, :self.modes2, :self.modes3], self.weights2)
        
        # 逆傅里叶变换
        x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)), dim=[-3, -2, -1])
        
        return x
    
    def complex_mul3d(self, input, weights):
        """3D复数乘法"""
        return torch.einsum("bixyz,ioxyz->boxyz", input, weights)


class FNO3d(nn.Module):
    """3D傅里叶神经算子，用于三维电磁场预测"""
    
    def __init__(
        self,
        modes1=8,
        modes2=8,
        modes3=8,
        width=64,
        n_layers=4,
        in_dim=7,
        out_dim=4,
        activation='gelu',
        pad_x=5,
        pad_y=0,
        pad_z=0
    ):
        super(FNO3d, self).__init__()
        """初始化3D傅里叶神经算子
        
        参数:
            modes1: x方向傅里叶模式数
            modes2: y方向傅里叶模式数
            modes3: z方向傅里叶模式数
            width: 隐藏层宽度
            n_layers: FNO层数
            in_dim: 输入维度 (通常为7: t, x, y, z, u, v, w)
            out_dim: 输出维度 (通常为4: u, v, Bx, By)
            activation: 激活函数
            pad_x: x方向填充
            pad_y: y方向填充
            pad_z: z方向填充
        """
        self.modes1 = modes1
        self.modes2 = modes2
        self.modes3 = modes3
        self.width = width
        self.n_layers = n_layers
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.pad_x = pad_x
        self.pad_y = pad_y
        self.pad_z = pad_z
        
        # 输入投影层
        self.fc0 = nn.Linear(in_dim, self.width)
        
        # 傅里叶层
        self.sp_convs = nn.ModuleList()
        for _ in range(self.n_layers):
            self.sp_convs.append(
                SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
            )
        
        # 卷积层
        self.convs = nn.ModuleList()
        for _ in range(self.n_layers):
            self.convs.append(nn.Conv3d(self.width, self.width, 1))
        
        # 激活函数
        if activation == 'gelu':
            self.activation = F.gelu
        elif activation == 'relu':
            self.activation = F.relu
        else:
            self.activation = F.gelu
        
        # 输出投影层
        self.fc1 = nn.Linear(self.width, 128)
        self.fc2 = nn.Linear(128, self.out_dim)
    
    def forward(self, x):
        """前向传播
        
        参数:
            x: 输入张量，形状为 [batch_size, x_dim, y_dim, z_dim, in_dim]
            
        返回:
            输出张量，形状为 [batch_size, x_dim, y_dim, z_dim, out_dim]
        """
        # 输入投影
        x = self.fc0(x)
        x = x.permute(0, 4, 1, 2, 3)  # [batch_size, width, x_dim, y_dim, z_dim]
        
        # 填充
        if self.pad_x > 0 or self.pad_y > 0 or self.pad_z > 0:
            x = F.pad(x, [0, self.pad_z, 0, self.pad_y, 0, self.pad_x])
        
        # FNO层
        for i in range(self.n_layers):
            x1 = self.sp_convs[i](x)
            x2 = self.convs[i](x)
            x = x1 + x2
            x = self.activation(x)
        
        # 去除填充
        if self.pad_x > 0 or self.pad_y > 0 or self.pad_z > 0:
            if self.pad_z > 0:
                x = x[..., :-self.pad_z]
            if self.pad_y > 0:
                x = x[..., :-self.pad_y, :]
            if self.pad_x > 0:
                x = x[..., :-self.pad_x, :, :]
        
        # 输出投影
        x = x.permute(0, 2, 3, 4, 1)  # [batch_size, x_dim, y_dim, z_dim, width]
        x = self.fc1(x)
        x = self.activation(x)
        x = self.fc2(x)
        
        return x


class EMFieldPredictor(nn.Module):
    """三维电磁场预测模型"""
    
    def __init__(
        self,
        model_params=None,
        device='cuda' if torch.cuda.is_available() else 'cpu'
    ):
        super(EMFieldPredictor, self).__init__()
        """初始化电磁场预测模型
        
        参数:
            model_params: 模型参数字典
            device: 计算设备
        """
        self.device = device
        
        # 默认参数
        default_params = {
            'modes1': 8,
            'modes2': 8,
            'modes3': 8,
            'width': 64,
            'n_layers': 4,
            'in_dim': 7,  # t, x, y, z, u, v, w
            'out_dim': 4,  # u, v, Bx, By
            'activation': 'gelu',
            'pad_x': 5,
            'pad_y': 0,
            'pad_z': 0
        }
        
        # 更新参数
        if model_params is not None:
            default_params.update(model_params)
        
        self.model_params = default_params
        
        # 创建FNO模型
        self.fno = FNO3d(**self.model_params).to(self.device)
        
        # 输入输出归一化参数
        self.input_norm = torch.tensor(self.model_params.get('input_norm', [1.0]*7)).to(self.device)
        self.output_norm = torch.tensor(self.model_params.get('output_norm', [1.0]*4)).to(self.device)
    
    def forward(self, x):
        """前向传播
        
        参数:
            x: 输入张量，形状为 [batch_size, x_dim, y_dim, z_dim, in_dim]
            
        返回:
            输出张量，形状为 [batch_size, x_dim, y_dim, z_dim, out_dim]
        """
        # 输入归一化
        x = x / self.input_norm
        
        # 模型预测
        output = self.fno(x)
        
        # 输出反归一化
        output = output * self.output_norm
        
        return output
    
    def predict(self, x):
        """预测函数
        
        参数:
            x: 输入张量，形状为 [batch_size, x_dim, y_dim, z_dim, in_dim]
            
        返回:
            输出张量，形状为 [batch_size, x_dim, y_dim, z_dim, out_dim]
        """
        self.eval()
        with torch.no_grad():
            return self.forward(x)
    
    def save_checkpoint(self, filepath):
        """保存模型检查点
        
        参数:
            filepath: 保存路径
        """
        torch.save({
            'model_state_dict': self.state_dict(),
            'model_params': self.model_params,
            'input_norm': self.input_norm.cpu().numpy(),
            'output_norm': self.output_norm.cpu().numpy()
        }, filepath)
    
    def load_checkpoint(self, filepath):
        """加载模型检查点
        
        参数:
            filepath: 加载路径
        """
        checkpoint = torch.load(filepath, map_location=self.device)
        self.load_state_dict(checkpoint['model_state_dict'])
        self.model_params = checkpoint['model_params']
        self.input_norm = torch.tensor(checkpoint['input_norm']).to(self.device)
        self.output_norm = torch.tensor(checkpoint['output_norm']).to(self.device)